#include "Compiler.h"
#include "AST.h"

#include <iostream>
#include <fstream>
#include <cstdlib>
#include <stack>
#include <filesystem>

std::unordered_map<std::string, std::string> long_option_to_short_option = {
	{ "--help", "-h" }, 
	{ "--purify", "-p" }, 
	{ "--assemble", "-a" }, 
	{ "--execute", "-e" }, 
	{ "--file", "-f" }, 
	{ "--capacity", "-c" }, 
	{ "--intermediate-dir", "-i" }, 
	{ "--target-dir", "-t" }
};

nmspc::Compiler::Compiler() :paper_tape_capacity(65536), init_ch(0)
{
}

std::string nmspc::Compiler::getFileExtension(const std::string& file_name)
{
	std::size_t dot_position = file_name.find_last_of('.');
	if (dot_position != std::string::npos) 
	{
		return file_name.substr(dot_position + 1);
	}
	return "";
}

std::string nmspc::Compiler::readTextFile(const std::string& file_name)
{
	std::ifstream fin(file_name.c_str(), std::ios::in);
	std::string full_code;
	if (fin.is_open())
	{
		std::string assist;
		while (std::getline(fin, assist))
		{
			full_code += std::move(assist);
		}
		fin.close();
	}
	else
	{
		throw std::runtime_error("File \"" + file_name + "\" does not exist!");
	}
	return full_code;
}

void nmspc::Compiler::writeTextFile(const std::string& file_name, const std::string& text)
{
	std::ofstream fout(file_name.c_str(), std::ios::out);
	std::string full_code;
	if (fout.is_open())
	{
		fout << text << std::endl;
		fout.close();
	}
	else
	{
		throw std::runtime_error("File \"" + file_name + "\" does not exist!");
	}
}

std::vector<std::string> nmspc::Compiler::split(const std::string& cmd, const char& skip)
{
	std::vector<std::string> ans;
	std::string assist;
	for (auto& ch : cmd)
	{
		if (ch == skip && !assist.empty())
		{
			ans.push_back(assist);
			assist.clear();
		}
		else
		{
			assist.push_back(ch);
		}
	}
	if (!assist.empty())ans.push_back(assist);
	return ans;
}

std::unordered_map<std::string, std::vector<std::string>> nmspc::Compiler::parseCommand(const std::vector<std::string>& splited_command)
{
	std::unordered_map<std::string, std::vector<std::string>> cmd_value_map;
	std::string cur_option = "--commandName";
	for (std::int32_t i = 0; i < splited_command.size(); i++)
	{
		if (splited_command[i].size() > 1 && splited_command[i][0] == '-' && !(splited_command[i][1] >= '0' && splited_command[i][1] <= '9'))
		{
			if (long_option_to_short_option.find(splited_command[i]) != long_option_to_short_option.end())
			{
				cur_option = long_option_to_short_option[splited_command[i]];
			}
			else
			{
				cur_option = splited_command[i];
			}
			if (cmd_value_map.find(cur_option) == cmd_value_map.end())
			{
				cmd_value_map[cur_option] = {};
			}
		}
		else if (!cur_option.empty())
		{
			cmd_value_map[cur_option].push_back(splited_command[i]);
		}
	}
	return cmd_value_map;
}

std::string nmspc::Compiler::purify(const std::string& full_code)
{
	std::string pure_brain_fuck_code;
	for (const auto& ch : full_code)
	{
		if (ch == '+' ||
			ch == '-' ||
			ch == '<' ||
			ch == '>' ||
			ch == '[' ||
			ch == ']' ||
			ch == ',' ||
			ch == '.')
		{
			pure_brain_fuck_code.push_back(ch);
		}
	}
	return pure_brain_fuck_code;
}

std::string nmspc::Compiler::assemble(const std::string& pure_brain_fuck_code)
{
	std::int64_t pbf_index = 0;
	std::uint64_t pbf_size = pure_brain_fuck_code.size();
	nmspc::LoopASTNode* main_ast_node = nmspc::generateLoopASTNode(pure_brain_fuck_code, 0, pure_brain_fuck_code.size() - 1);
	std::vector<nmspc::ASTNode*> ast_nodes;
	std::string assemble_code;
	assemble_code += "section .bss\n";
	assemble_code += "    paper_tape resb " + std::to_string(this->paper_tape_capacity) + "\n";
	assemble_code += "    paper_tape_ptr resq 1\n";
	assemble_code += "section .data\n";
	assemble_code += "    fmt db \"%c\", 0\n";
	assemble_code += "    paper_tape_capacity dq " + std::to_string(this->paper_tape_capacity) + "\n";
	assemble_code += "section .text\n";
	assemble_code += "    global main\n";
	assemble_code += "    extern printf\n";
	assemble_code += "    extern _getch\n";
	assemble_code += "    extern ExitProcess\n";
	assemble_code += "main:\n";
	assemble_code += "    lea rax, [rel paper_tape]\n";
	assemble_code += "    mov [rel paper_tape_ptr], rax\n";
	assemble_code += "    mov rcx, " + std::to_string(this->paper_tape_capacity) + "\n";
	assemble_code += "    mov rbx, 0\n";
	assemble_code += "init_tape:\n";
	assemble_code += "    mov byte [rax + rbx], " + std::to_string(static_cast<std::int32_t>(this->init_ch)) + "\n";
	assemble_code += "    inc rbx\n";
	assemble_code += "    loop init_tape\n";
	assemble_code += main_ast_node->generateAssembleCode(0);
	assemble_code += "    xor rcx, rcx\n";
	assemble_code += "    call ExitProcess";
	return assemble_code;
}

void nmspc::Compiler::execute(const std::string& asm_file_name, const std::string& object_file_name, const std::string& execute_file_name)
{
	std::filesystem::path tool_dir{ ".\\vendor\\bin\\" };
	std::string tool_dir_str = std::filesystem::absolute(tool_dir).string();
	std::filesystem::path dll_dir{ ".\\vendor\\lib" };
	std::string dll_dir_str = std::filesystem::absolute(dll_dir).string();
	std::string nasm_command = tool_dir_str + "nasm.exe -f win64 -o " + object_file_name + " " + asm_file_name;
	if (std::system(nasm_command.c_str()) != 0)
	{
		throw std::runtime_error("NASM failed to assemble the code.");
	}
	std::string ld_command = tool_dir_str + "ld.exe -o " + execute_file_name + " " + object_file_name + " -e main -L" + dll_dir_str + " -lmsvcrt -lkernel32 --verbose -v";
	if (std::system(ld_command.c_str()) != 0)
	{
		throw std::runtime_error("LD failed to link the object file.");
	}
}

std::int32_t nmspc::Compiler::entryPoint(int argc, char* argv[])
{
	std::vector<std::string> splited_command(argc);
	for (int i = 0; i < argc; i++)
	{
		splited_command[i] = std::move(std::string(argv[i]));
	}
	std::unordered_map<std::string, std::vector<std::string>> cmd_value_map = std::move(this->parseCommand(splited_command));
	if (cmd_value_map.find("-h") != cmd_value_map.end() || cmd_value_map.find("--help") != cmd_value_map.end())
	{
		std::cerr << "\033[36m" << "Welcome to the BrainFuck compiler!" << "\033[0m" << std::endl;
		std::cerr << "\033[32m" << "usage; " << "\033[0m" << std::endl;
		std::cerr << "  BrainFuck [[option] [value]]..." << std::endl;
		std::cerr << "\033[32m" << "options: " << "\033[0m" << std::endl;
		std::cerr << "  -a, --assemble            // Spycify the assemble file's name." << std::endl;
		std::cerr << "  -c, --capacity            // Spycify the capacity of the paper tape." << std::endl;
		std::cerr << "  -e, --execute             // Spycify the executable file's name." << std::endl;
		std::cerr << "  -f, --file                // Spycify the source code file name." << std::endl;
		std::cerr << "  -h, --help                // Display help document." << std::endl;
		std::cerr << "  -i, --intermediate-dir    // Spycify the capacity of the paper tape." << std::endl;
		std::cerr << "      --init-paper-tape-ch  // Spycify the init-ch of the paper tape." << std::endl;
		std::cerr << "  -o, --object              // Spycify the object file's name." << std::endl;
		std::cerr << "  -p, --purify              // Spycify the pure-brain-fuck file's name." << std::endl;
		std::cerr << "  -t, --target-dir          // Spycify the capacity of the paper tape." << std::endl;
		return 0;
	}
	if (cmd_value_map.find("-c") != cmd_value_map.end() && !cmd_value_map["-c"].empty())
	{
		try
		{
			this->paper_tape_capacity = std::stoull(cmd_value_map["-c"][0]);
		}
		catch (const std::exception& e)
		{
			std::cerr << "\033[31m" << e.what() << "\033[0m" << std::endl;
		}
	}
	if (cmd_value_map.find("--init-paper-tape-ch") != cmd_value_map.end() && !cmd_value_map["--init-paper-tape-ch"].empty() && !cmd_value_map["--init-paper-tape-ch"][0].empty())
	{
		std::string cur_init_paper_tape_ch = cmd_value_map["--init-paper-tape-ch"][0];
		if (cur_init_paper_tape_ch[0] == 'n' && cur_init_paper_tape_ch.size() > 1)
		{
			cur_init_paper_tape_ch.erase(cur_init_paper_tape_ch.begin());
			try
			{
				this->init_ch = std::stoi(cur_init_paper_tape_ch);
			}
			catch (const std::exception& e)
			{
				std::cerr << "\033[31m" << "Error: " << e.what() << "\033[0m" << std::endl;
				return 1;
			}
		}
		else
		{
			this->init_ch = cmd_value_map["--init-paper-tape-ch"][0][0];
		}
	}
	if (cmd_value_map.find("-f") != cmd_value_map.end())
	{
		if (cmd_value_map["-f"].size() < 1)
		{
			std::cerr << "\033[31m" << "The option \"-f\" should have at least one value." << "\033[0m" << std::endl;
			return 1;
		}
		std::vector<std::string> src_code_full_file_name_list = cmd_value_map["-f"];
		for (std::string src_code_full_file_name : src_code_full_file_name_list)
		{
			std::filesystem::path src_code_file_path(src_code_full_file_name);
			std::string src_code_file_extension_name = src_code_file_path.extension().string();
			std::string src_code_file_name = src_code_file_path.stem().string();
			std::string intermediate_dir = std::filesystem::absolute(src_code_file_path).parent_path().string() + "\\";
			std::string target_dir = std::filesystem::absolute(src_code_file_path).parent_path().string() + "\\";
			if (cmd_value_map.find("-i") != cmd_value_map.end() && !cmd_value_map["-i"].empty())
			{
				intermediate_dir = cmd_value_map["-i"][0];
				if (intermediate_dir.back() != '\\' && intermediate_dir.back() != '/')
				{
					intermediate_dir += '/';
				}
				try 
				{
					if (!std::filesystem::exists(intermediate_dir))
					{
						std::filesystem::create_directories(intermediate_dir);
					}
				}
				catch (const std::filesystem::filesystem_error& e)
				{
					std::cerr << "\033[31m" << "Error: " << e.what() << "\033[0m" << std::endl;
				}
			}
			if (cmd_value_map.find("-t") != cmd_value_map.end() && !cmd_value_map["-t"].empty())
			{
				target_dir = cmd_value_map["-t"][0];
				if (target_dir.back() != '\\' && target_dir.back() != '/')
				{
					target_dir += '/';
				}
				try
				{
					if (!std::filesystem::exists(target_dir))
					{
						std::filesystem::create_directories(target_dir);
					}
				}
				catch (const std::filesystem::filesystem_error& e)
				{
					std::cerr << "\033[31m" << "Error: " << e.what() << "\033[0m" << std::endl;
				}
			}
			if (!src_code_file_extension_name.empty())
			{
				for (auto& ch : src_code_file_extension_name)
				{
					if (ch >= 'A' && ch <= 'Z')ch += 32;
				}
				if (src_code_file_extension_name != ".bf" && src_code_file_extension_name != ".brainfuck" && src_code_file_extension_name != ".pbf")
				{
					std::cerr << "\033[31m" << "The file \"" << src_code_full_file_name << "\" has an extension of \"" << src_code_file_extension_name << "\" and cannot be compiled!" << "\033[0m" << std::endl;
					continue;
				}
				std::string pure_brain_fuck_code;
				if (src_code_file_extension_name == ".bf" || src_code_file_extension_name == ".brainfuck")
				{
					std::string full_code;
					try
					{
						full_code = this->readTextFile(src_code_full_file_name);
					}
					catch (const std::exception& e)
					{
						std::cerr << "\033[31m" << "Error: " << e.what() << "\033[0m" << std::endl;
						continue;
					}
					pure_brain_fuck_code = this->purify(full_code);
				}
				if (cmd_value_map.find("-p") != cmd_value_map.end())
				{
					std::string purify_file_name = std::move(intermediate_dir + src_code_file_name);
					if (cmd_value_map["-p"].size() == 1)
					{
						purify_file_name = std::move(intermediate_dir + cmd_value_map["-p"][0]);
					}
					else if (cmd_value_map["-p"].size() > 1)
					{
						std::cerr << "\033[31m" << "The option \"-p\" only accepts one value!" << "\033[0m" << std::endl;
						continue;
					}
					purify_file_name += ".pbf";
					try
					{
						this->writeTextFile(purify_file_name, pure_brain_fuck_code);
					}
					catch (const std::exception& e)
					{
						std::cerr << "\033[31m" << "Error: " << e.what() << "\033[0m" << std::endl;
						continue;
					}
				}
				if (src_code_file_extension_name == ".pbf")
				{
					try
					{
						pure_brain_fuck_code = this->readTextFile(src_code_full_file_name);
					}
					catch (const std::exception& e)
					{
						std::cerr << "\033[31m" << "Error: " << e.what() << "\033[0m" << std::endl;
						continue;
					}
				}
				std::string assemble_code;
				assemble_code = this->assemble(pure_brain_fuck_code);
				std::string assemble_file_name = std::move(intermediate_dir + src_code_file_name);
				if (cmd_value_map.find("-a") != cmd_value_map.end())
				{
					if (cmd_value_map["-a"].size() == 1)
					{
						assemble_file_name = std::move(intermediate_dir + cmd_value_map["-a"][0]);
					}
					else if (cmd_value_map["-a"].size() > 1)
					{
						std::cerr << "\033[31m" << "The option \"-a\" only accepts one value" << "\033[0m" << std::endl;
						continue;
					}
				}
				assemble_file_name += ".asm";
				try
				{
					this->writeTextFile(assemble_file_name, assemble_code);
				}
				catch (const std::exception& e)
				{
					std::cerr << "\033[31m" << "Error: " << e.what() << "\033[0m" << std::endl;
					continue;
				}
				std::string execute_file_name = std::move(target_dir + src_code_file_name);
				if (cmd_value_map.find("-e") != cmd_value_map.end())
				{
					if (cmd_value_map["-e"].size() == 1)
					{
						execute_file_name = std::move(target_dir + cmd_value_map["-e"][0]);
					}
					else if (cmd_value_map["-e"].size() > 1)
					{
						std::cerr << "\033[31m" << "The option \"-e\" only accepts one value" << "\033[0m" << std::endl;
						continue;
					}
				}
				execute_file_name += ".exe";
				std::string object_file_name = std::move(intermediate_dir + src_code_file_name);
				if (cmd_value_map.find("-o") != cmd_value_map.end())
				{
					if (cmd_value_map["-o"].size() == 1)
					{
						object_file_name = std::move(intermediate_dir + cmd_value_map["-o"][0]);
					}
					else if (cmd_value_map["-o"].size() > 1)
					{
						std::cerr << "\033[31m" << "The option \"-o\" only accepts one value" << "\033[0m" << std::endl;
						continue;
					}
				}
				object_file_name += ".obj";
				try
				{
					this->execute(assemble_file_name, object_file_name, execute_file_name);
				}
				catch (const std::exception& e)
				{
					std::cerr << "\033[31m" << "Error: " << e.what() << "\033[0m" << std::endl;
				}
			}
			else
			{
				std::cerr << "\033[31m" << "Please spycify the source code file extension name!" << "\033[0m" << std::endl;
				return 1;
			}
		}
	}
	else
	{
		std::cerr << "\033[31m" << "Please spycify the source code file name!" << "\033[0m" << std::endl;
		return 1;
	}
	return 0;
}
